#include "nm_dir.h"

kdirectory::kdirectory()
{
	m_filehandle = 0;
	m_fileobj = NULL;
}

void kdirectory::Release()
{
	Close();
}

VOID kdirectory::Close()
{
	if (m_fileobj != NULL) {

		ObDereferenceObject(m_fileobj);
		m_fileobj = NULL;
	}

	if (m_filehandle != NULL) {

		ZwClose(m_filehandle);
		m_filehandle = 0;
	}
}

NTSTATUS kdirectory::Create( PWCHAR FileName, ULONG CreateDisposition /*= FILE_OPEN*/, ACCESS_MASK DesiredAccess /*= GENERIC_READ*/, ULONG ShareAccess /*= FILE_SHARE_READ */ )
{
	UNICODE_STRING Name;

	RtlInitUnicodeString(&Name, FileName);

	return Create(&Name, CreateDisposition, DesiredAccess, ShareAccess);
}

NTSTATUS kdirectory::Create( PUNICODE_STRING FileName, ULONG CreateDisposition /*= FILE_OPEN*/, ACCESS_MASK DesiredAccess /*= GENERIC_READ*/, ULONG ShareAccess /*= FILE_SHARE_READ */ )
{
	NTSTATUS ns;
	OBJECT_ATTRIBUTES oa;
	IO_STATUS_BLOCK Iob;

	Close();

	InitializeObjectAttributes(&oa, FileName, OBJ_KERNEL_HANDLE | OBJ_CASE_INSENSITIVE, 0, NULL);

	ns = IoCreateFile(&m_filehandle, 
		DesiredAccess | SYNCHRONIZE, 
		&oa, 
		&Iob, 
		NULL, 
		FILE_ATTRIBUTE_NORMAL, 
		ShareAccess, 
		CreateDisposition, 
		FILE_DIRECTORY_FILE | FILE_SYNCHRONOUS_IO_NONALERT, 
		NULL, 
		0, 
		CreateFileTypeNone, 
		NULL, 
		IO_NO_PARAMETER_CHECKING);

	if (!NT_SUCCESS(ns)) {

		return ns;
	}

	ns = ObReferenceObjectByHandle(m_filehandle, 0, *IoFileObjectType, KernelMode, (PVOID *)&m_fileobj, NULL);
	if (!NT_SUCCESS(ns)) {

		ZwClose(m_filehandle);
		m_filehandle = 0;
	}

	return ns;
}

NTSTATUS kdirectory::Create( kstringw& FileName, ULONG CreateDisposition /*= FILE_OPEN*/, ACCESS_MASK DesiredAccess /*= GENERIC_READ*/, ULONG ShareAccess /*= FILE_SHARE_READ */ )
{
	return Create((PUNICODE_STRING)FileName, CreateDisposition, DesiredAccess, ShareAccess);
}

NTSTATUS kdirectory::Delete( BOOLEAN Del /*= TRUE*/ )
{
	NTSTATUS ns;
	IO_STATUS_BLOCK Iob;
	FILE_DISPOSITION_INFORMATION FileDispInfo;

	FileDispInfo.DeleteFile = Del;

	ns = ZwSetInformationFile(m_filehandle,
		&Iob,
		&FileDispInfo,
		sizeof(FileDispInfo),
		FileDispositionInformation);

	return ns;
}

HANDLE kdirectory::GetHandle()
{
	return m_filehandle;
}

NTSTATUS kdirectory::GetName(kstringw& String)
{
	kstringw Path;
	NTSTATUS ns;
	INT Pos;
	ULONG RightCch;
	ns = GetPath(Path);
	if (!NT_SUCCESS(ns)) {

		return ns;
	}

	Pos = Path.ReverseFind(L'\\');
	if (Pos == -1) {

		return STATUS_UNSUCCESSFUL;
	}

	Pos++;

	RightCch = Path.GetLength() - Pos;

	Path.Right(RightCch);

	return ns;
}

PFILE_OBJECT kdirectory::GetObject()
{
	return m_fileobj;
}

NTSTATUS kdirectory::GetPath(kstringw& String)
{
	PUCHAR Buffer;
	kstringw Path;
	NTSTATUS ns;

	Buffer = new UCHAR[1024];
	RtlZeroMemory(Buffer, 1024);

	ns = GetPath((PFILE_NAME_INFORMATION)Buffer, 1024 - sizeof(WCHAR));
	if (!NT_SUCCESS(ns)) {

		delete[] Buffer;
		return ns;
	}

	String.Copy(((PFILE_NAME_INFORMATION)Buffer)->FileName, 
		((PFILE_NAME_INFORMATION)Buffer)->FileNameLength / sizeof(WCHAR));

	delete[] Buffer;
	return ns;
}

NTSTATUS kdirectory::GetPath( PFILE_NAME_INFORMATION NameInfo, ULONG Length)
{
	NTSTATUS ns;
	IO_STATUS_BLOCK Iob;

	return ZwQueryInformationFile(m_filehandle,
		&Iob,
		&NameInfo,
		Length,
		FileNameInformation);
}

NTSTATUS kdirectory::MakeDir( PUNICODE_STRING FileName )
{
	NTSTATUS ns;
	PWCHAR p = FileName->Buffer;
	USHORT i = 0;
	PWCHAR q;
	kstringw dir;

	ns = Create(FileName, FILE_OPEN_IF, GENERIC_READ | GENERIC_WRITE);
	if (NT_SUCCESS(ns)) {

		Close();
		return ns;
	}

	q = dir.Allocate(FileName->Length / sizeof(WCHAR));
	while (i < FileName->Length) {

		if (L'\\' == *p) {

			if (L':' != *(p - 1)) {
				Create(dir.GetUnicodeString(), FILE_OPEN_IF, GENERIC_READ | GENERIC_WRITE);
			}
		}
		*q++ = *p++;
		*q = L'\0';
		i += 2;
	}

	ns = Create(dir.GetUnicodeString(), FILE_OPEN_IF, GENERIC_READ | GENERIC_WRITE);
	Close();

	return ns;
}

NTSTATUS kdirectory::MakeDir( kstringw& FileName )
{
	return MakeDir(FileName.GetUnicodeString());
}

NTSTATUS kdirectory::MakeDir( PWCHAR FileName )
{
	UNICODE_STRING UniString;

	RtlInitUnicodeString(&UniString, FileName);

	return MakeDir(&UniString);
}


NTSTATUS kdirectory::Rename( PFILE_RENAME_INFORMATION RenameInfo, ULONG Length )
{
	IO_STATUS_BLOCK Iob;

	return ZwSetInformationFile(m_filehandle,
		&Iob,
		RenameInfo,
		Length,
		FileRenameInformation);
}

NTSTATUS kdirectory::Rename( kstringw& Name, BOOLEAN ReplaceIfExists )
{
	NTSTATUS ns;
	PUCHAR Buffer;
	PFILE_RENAME_INFORMATION RenameInfo;

	Buffer = new UCHAR[1024];
	RtlZeroMemory(Buffer, 1024);
	RenameInfo = (PFILE_RENAME_INFORMATION)Buffer;

	RenameInfo->ReplaceIfExists = ReplaceIfExists;
	RenameInfo->RootDirectory = NULL;
	RenameInfo->FileNameLength = Name.GetLength() * sizeof(WCHAR);
	RtlStringCbCopyW(RenameInfo->FileName, 1024 - sizeof(FILE_RENAME_INFORMATION), Name.GetString());

	ns = Rename(RenameInfo, 1024);

	delete[] Buffer;

	return ns;
}

NTSTATUS kdirectory::QueryDir( 
							  PVOID FileInformation, 
							  ULONG Length, 
							  PUNICODE_STRING FileName /*= NULL*/, 
							  FILE_INFORMATION_CLASS FileInformationClass /*= FileBothDirectoryInformation*/, 
							  BOOLEAN ReturnSingleEntry /*= TRUE */ 
							  )
{
	IO_STATUS_BLOCK Iob;
	return ZwQueryDirectoryFile(m_filehandle, 
		NULL, 
		NULL, 
		NULL, 
		&Iob, 
		FileInformation, 
		Length, 
		FileInformationClass, 
		ReturnSingleEntry, 
		FileName, 
		FALSE);
}

BOOLEAN kdirectory::IsDirectory( kstringw& String )
{
	return kdirectory::IsDirectory(String.GetUnicodeString());
}

BOOLEAN kdirectory::IsDirectory( PUNICODE_STRING UniString )
{
	NTSTATUS ns;
	OBJECT_ATTRIBUTES oa;
	IO_STATUS_BLOCK Iob;
	HANDLE FileHandle;
	FILE_STANDARD_INFORMATION FileInfo;

	InitializeObjectAttributes(&oa, UniString, OBJ_KERNEL_HANDLE | OBJ_CASE_INSENSITIVE, 0, NULL);

	ns = IoCreateFile(&FileHandle, 
		GENERIC_READ, 
		&oa, 
		&Iob, 
		NULL, 
		FILE_ATTRIBUTE_NORMAL, 
		FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE, 
		FILE_OPEN, 
		FILE_SYNCHRONOUS_IO_NONALERT, 
		NULL, 
		0, 
		CreateFileTypeNone, 
		NULL, 
		IO_NO_PARAMETER_CHECKING);

	if (!NT_SUCCESS(ns)) {

		return FALSE;
	}

	ns = ZwQueryInformationFile(FileHandle, &Iob, &FileInfo, sizeof(FileInfo), FileStandardInformation);

	ZwClose(FileHandle);
	if (!NT_SUCCESS(ns)) {

		return FALSE;
	}

	return FileInfo.Directory;

}

BOOLEAN kdirectory::IsDirectory( PWCHAR String )
{
	UNICODE_STRING UniString;

	RtlInitUnicodeString(&UniString, String);

	return kdirectory::IsDirectory(&UniString);
}